from typing import Literal

from pydantic import BaseModel, ConfigDict, Field

ModelTask = Literal["automatic-speech-recognition", "text-to-speech", "speaker-embedding", "voice-activity-detection"]

# https://platform.openai.com/docs/api-reference/audio/createSpeech#audio-createspeech-response_format
DEFAULT_SPEECH_RESPONSE_FORMAT = "mp3"

# https://platform.openai.com/docs/api-reference/audio/createSpeech#audio-createspeech-voice
# https://platform.openai.com/docs/guides/text-to-speech/voice-options
OPENAI_SUPPORTED_SPEECH_VOICE_NAMES = ("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse")

# https://platform.openai.com/docs/guides/text-to-speech/supported-output-formats
type SpeechResponseFormat = Literal["pcm", "mp3", "wav", "flac", "opus", "aac"]
SUPPORTED_SPEECH_RESPONSE_FORMATS = ("pcm", "mp3", "wav", "flac", "opus", "aac")

MIN_SPEECH_SAMPLE_RATE = 8000
MAX_SPEECH_SAMPLE_RATE = 48000


# https://github.com/openai/openai-openapi/blob/master/openapi.yaml#L11146
class Model(BaseModel):
    """There may be additional fields in the response that are specific to the model type."""

    id: str
    """The model identifier, which can be referenced in the API endpoints."""
    created: int = 0
    """The Unix timestamp (in seconds) when the model was created."""
    object: Literal["model"] = "model"
    """The object type, which is always "model"."""
    owned_by: str
    """The organization that owns the model."""
    language: list[str] | None = None
    """List of ISO 639-3 supported by the model. It's possible that the list will be empty. This field is not a part of the OpenAI API spec and is added for convenience."""

    task: ModelTask  # TODO: make a list?

    model_config = ConfigDict(extra="allow")


# https://github.com/openai/openai-openapi/blob/master/openapi.yaml#L8730
class ListModelsResponse(BaseModel):
    data: list[Model]
    object: Literal["list"] = "list"


# https://github.com/openai/openai-openapi/blob/master/openapi.yaml#L10909
TimestampGranularities = list[Literal["segment", "word"]]


DEFAULT_TIMESTAMP_GRANULARITIES: TimestampGranularities = ["segment"]
TIMESTAMP_GRANULARITIES_COMBINATIONS: list[TimestampGranularities] = [
    [],  # should be treated as ["segment"]. https://platform.openai.com/docs/api-reference/audio/createTranscription#audio-createtranscription-timestamp_granularities
    ["segment"],
    ["word"],
    ["word", "segment"],
    ["segment", "word"],  # same as ["word", "segment"] but order is different
]


class EmbeddingObject(BaseModel):
    object: Literal["embedding"] = "embedding"
    """The object type, which is always "embedding"."""
    index: Literal[0] = 0
    """The index (always 0) of the embedding in the list of embeddings."""
    embedding: list[float]
    """The embedding vector, which is a list of floats. The length of vector depends on the model used."""


class EmbeddingUsage(BaseModel):
    prompt_tokens: int
    """The number of tokens in the input audio."""
    total_tokens: int
    """The total number of tokens used (same as prompt_tokens for audio embeddings)."""


class CreateEmbeddingResponse(BaseModel):
    object: Literal["list"] = "list"
    """The object type, which is always "list"."""
    data: list[EmbeddingObject] = Field(..., min_length=1, max_length=1)
    """The list of embeddings generated by the model. Always contains exactly one embedding."""
    model: str
    """The model used to generate the embeddings."""
    usage: EmbeddingUsage
    """Usage statistics for the API call."""


# NOTE: I define these here because they aren't defined in the openai-python package. Once they are added there, we can remove these definitions.


class SpeechAudioDeltaEvent(BaseModel):
    """Emitted for each chunk of audio data generated during speech synthesis."""

    type: Literal["speech.audio.delta"] = "speech.audio.delta"
    """The type of the event. Always speech.audio.delta."""
    audio: str
    """A chunk of Base64-encoded audio data."""


class SpeechAudioTokenUsage(BaseModel):
    """Token usage statistics for the request."""

    input_tokens: int
    """Number of input tokens in the prompt."""
    output_tokens: int
    """Number of output tokens generated."""
    total_tokens: int
    """Total number of tokens used (input + output)."""


class SpeechAudioDoneEvent(BaseModel):
    """Emitted when the speech synthesis is complete and all audio has been streamed."""

    type: Literal["speech.audio.done"] = "speech.audio.done"
    """The type of the event. Always speech.audio.done."""
    token_usage: SpeechAudioTokenUsage
    """Token usage statistics for the request."""
